package woaini.wind.ai.controller;

import cn.hutool.core.collection.CollUtil;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.ollama.OllamaChatModel;
import dev.langchain4j.service.AiServices;
import lombok.RequiredArgsConstructor;
import org.springframework.stereotype.Service;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

@Service
@RequiredArgsConstructor
public class AssistantService {

  private static final OllamaChatModel ollamaChatModel;
  private final AssistantTools assistantTools;

  static {
    ollamaChatModel =
        OllamaChatModel.builder()
            .baseUrl("http://10.8.1.250:11434")
            .modelName("qwq:latest")
            .listeners(CollUtil.newArrayList(new MyChatModelListener()))
            .logRequests(true)
            .logResponses(true)
            .build();
  }

  public static final Map<Integer, MessageWindowChatMemory> MEMORY_MAP = new ConcurrentHashMap<>();

  public AIResponse assistant(String message, Integer memoryId, Boolean isNew) {

    // 记忆获取
    MessageWindowChatMemory messageWindowChatMemory = MEMORY_MAP.get(memoryId);

    if (messageWindowChatMemory == null) {
      messageWindowChatMemory =
          MessageWindowChatMemory.builder().id(memoryId).maxMessages(5).build();
      MEMORY_MAP.put(memoryId, messageWindowChatMemory);
    }
    if (isNew) {
      messageWindowChatMemory.clear();
    }

    MessageWindowChatMemory finalMessageWindowChatMemory = messageWindowChatMemory;
    Assistant assistant =
        AiServices.builder(Assistant.class)
            .chatLanguageModel(ollamaChatModel)
            .chatMemoryProvider(_ -> finalMessageWindowChatMemory)
            .tools(assistantTools)
            .build();

    String chat = assistant.chat(memoryId, message);

    return new AIResponse(chat);
  }
}
